
from torch.autograd import Variable

import model3

from SFNet import *

# 优化MDSF
class DMPHN_With_MDSF2(nn.Module):
    def __init__(self):
        super(DMPHN_With_MDSF2, self).__init__()
        self.encoder = nn.ModuleDict()
        self.decoder = nn.ModuleDict()
        for s in ['s1', 's2', 's3', 's4']:
            self.encoder[s] = nn.ModuleDict()
            self.decoder[s] = nn.ModuleDict()
            for lv in ['lv1', 'lv2', 'lv3']:
                self.encoder[s][lv] = model3.Encoder()
                self.decoder[s][lv] = model3.Decoder()

        

    def forward(self, inputs):

        images = {}
        feature = {}
        residual = {}
        for s in ['s1', 's2', 's3', 's4']:
            feature[s] = {}  # encoder output
            residual[s] = {}  # decoder output

        images['lv1'] = Variable(inputs - 0.5)
        H = images['lv1'].size(2)
        W = images['lv1'].size(3)

        images['lv2_1'] = images['lv1'][:, :, 0:int(H / 2), :]
        images['lv2_2'] = images['lv1'][:, :, int(H / 2):H, :]
        images['lv3_1'] = images['lv2_1'][:, :, :, 0:int(W / 2)]
        images['lv3_2'] = images['lv2_1'][:, :, :, int(W / 2):W]
        images['lv3_3'] = images['lv2_2'][:, :, :, 0:int(W / 2)]
        images['lv3_4'] = images['lv2_2'][:, :, :, int(W / 2):W]

        for s, ps in zip(['s1', 's2', 's3', 's4'], [None, 's1', 's2', 's3']):
            if ps is None:
                feature[s]['lv3_1'] = self.encoder[s]['lv3'](images['lv3_1'])
                feature[s]['lv3_2'] = self.encoder[s]['lv3'](images['lv3_2'])
                feature[s]['lv3_3'] = self.encoder[s]['lv3'](images['lv3_3'])
                feature[s]['lv3_4'] = self.encoder[s]['lv3'](images['lv3_4'])
            else:
                feature[s]['lv3_1'] = self.encoder[s]['lv3'](images['lv3_1'] + residual[ps]['lv1'][:, :, 0:int(H / 2), 0:int(W / 2)])
                feature[s]['lv3_2'] = self.encoder[s]['lv3'](images['lv3_2'] + residual[ps]['lv1'][:, :, 0:int(H / 2), int(W / 2):W])
                feature[s]['lv3_3'] = self.encoder[s]['lv3'](images['lv3_3'] + residual[ps]['lv1'][:, :, int(H / 2):H, 0:int(W / 2)])
                feature[s]['lv3_4'] = self.encoder[s]['lv3'](images['lv3_4'] + residual[ps]['lv1'][:, :, int(H / 2):H, int(W / 2):W])

            feature[s]['lv3_top'] = torch.cat((feature[s]['lv3_1'], feature[s]['lv3_2']), 3)
            feature[s]['lv3_bot'] = torch.cat((feature[s]['lv3_3'], feature[s]['lv3_4']), 3)

            if ps is not None:
                feature[s]['lv3_top'] += feature[ps]['lv3_top']
                feature[s]['lv3_bot'] += feature[ps]['lv3_bot']

            residual[s]['lv3_top'] = self.decoder[s]['lv3'](feature[s]['lv3_top'])
            residual[s]['lv3_bot'] = self.decoder[s]['lv3'](feature[s]['lv3_bot'])

            ########################################

            if ps is None:
                feature[s]['lv2_1'] = self.encoder[s]['lv2'](images['lv2_1'] + residual[s]['lv3_top'])
                feature[s]['lv2_2'] = self.encoder[s]['lv2'](images['lv2_2'] + residual[s]['lv3_bot'])
            else:
                feature[s]['lv2_1'] = self.encoder[s]['lv2'](images['lv2_1'] + residual[s]['lv3_top'] + residual[ps]['lv1'][:, :, 0:int(H / 2), :])
                feature[s]['lv2_2'] = self.encoder[s]['lv2'](images['lv2_2'] + residual[s]['lv3_bot'] + residual[ps]['lv1'][:, :, int(H / 2):H, :])

            feature[s]['lv2_1'] += feature[s]['lv3_top']
            feature[s]['lv2_2'] += feature[s]['lv3_bot']

            if ps is not None:
                feature[s]['lv2_1'] += feature[ps]['lv2_1']
                feature[s]['lv2_2'] += feature[ps]['lv2_2']

            feature[s]['lv2'] = torch.cat((feature[s]['lv2_1'], feature[s]['lv2_2']), 2)

            residual[s]['lv2'] = self.decoder[s]['lv2'](feature[s]['lv2'])

            ########################################

            if ps is None:
                feature[s]['lv1'] = self.encoder[s]['lv1'](images['lv1'] + residual[s]['lv2'])
            else:
                feature[s]['lv1'] = self.encoder[s]['lv1'](images['lv1'] + residual[s]['lv2'] + residual[ps]['lv1'])

            feature[s]['lv1'] += feature[s]['lv2']

            if ps is not None:
                feature[s]['lv1'] += feature[ps]['lv1']

            residual[s]['lv1'] = self.decoder[s]['lv1'](feature[s]['lv1'])

        return residual['s1']['lv1'], residual['s2']['lv1'], residual['s3']['lv1'], residual['s4']['lv1']